import pygame
from pygame.locals import *
import math
from PIL import Image

ButtonPath = "./g_down.png"
TargetPath = "./simple1.png"
ButtonLength = 50
WindowLength = 800
WindowHeight = 600


class Brush():
    def __init__(self, screen):
        self.screen = screen
        self.color = (0, 0, 0)
        self.size = 12
        self.drawing = False
        self.last_pos = None
        self.space = 1
        self.style = False

    def start_draw(self, pos):
        self.drawing = True
        self.last_pos = pos

    def end_draw(self):
        self.drawing = False

    def draw(self, pos):
        if self.drawing:
            for p in self._get_points(pos):
                if self.style == False:
                    pygame.draw.circle(self.screen, self.color, p, self.size)
                else:
                    self.screen.blit(self.brush_now, p)
            self.last_pos = pos

    def _get_points(self, pos):
        points = [(self.last_pos[0], self.last_pos[1])]
        len_x = pos[0] - self.last_pos[0]
        len_y = pos[1] - self.last_pos[1]
        length = math.sqrt(len_x ** 2 + len_y ** 2)
        step_x = len_x / length
        step_y = len_y / length
        for i in range(int(length)):
            points.append((points[-1][0] + step_x, points[-1][1] + step_y))
        points = map(lambda x: (int(0.5 + x[0]), int(0.5 + x[1])), points)
        return list(set(points))


class Painter():
    def __init__(self):
        self.screen = pygame.display.set_mode((WindowLength, WindowHeight))
        pygame.display.set_caption("Painter")
        self.clock = pygame.time.Clock()
        self.brush = Brush(self.screen)
        self.menu = Menu(self.screen)

    def run(self):
        self.screen.fill((255, 255, 255))
        while True:
            self.clock.tick(30)
            for event in pygame.event.get():
                if event.type == QUIT:
                    return
                elif event.type == KEYDOWN:
                    if event.key == K_ESCAPE:
                        self.screen.fill((255, 255, 255))
                elif event.type == MOUSEBUTTONDOWN:
                    if ((event.pos)[0] <= 74 and self.menu.click_button(event.pos)):
                        pass
                    else:
                        self.brush.start_draw(event.pos)
                elif event.type == MOUSEMOTION:
                    self.brush.draw(event.pos)
                elif event.type == MOUSEBUTTONUP:
                    self.brush.end_draw()

            self.menu.draw()
            pygame.display.update()


class Menu():
    def __init__(self, screen):
        self.screen = screen
        self.brush = None

        self.buttons = [
            pygame.image.load(ButtonPath).convert_alpha()
        ]
        self.buttons_rect = []
        for (i, img) in enumerate(self.buttons):
            rect = pygame.Rect(10 + i % 2 * 32, 254 + i / 2 * 32, 32, 32)
            self.buttons_rect.append(rect)

    def draw(self):
        for (i, img) in enumerate(self.buttons):
            self.screen.blit(img, self.buttons_rect[i].topleft)

            # 点击判断部分

    def click_button(self, pos):
        for (i, rect) in enumerate(self.buttons_rect):
            if rect.collidepoint(pos):
                print("OK")
                pygame.image.save(self.screen, TargetPath)
                im = Image.open(TargetPath)
                cropped = im.crop((ButtonLength, 0, WindowHeight, WindowHeight))
                cropped.save(TargetPath)
                return True

    def set_brush(self, brush):
        self.brush = brush


if __name__ == '__main__':
    app = Painter()
    app.run()
